#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation

# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

from scipy.linalg import norm
from scipy import linspace, clip, exp

from utils.FemBases import LagrangeP1, LagrangeP2
from input_output.DerivativePlots import DerivativePlots

import models.NumericalOptionGBM as nOp
from models.AnalyticalFormulas import vanillaCallPrice, \
    vanillaPutPrice, cashOrNothingCall

from solvers.EuropeanTimeStepper import EuropeanTimeStepper
from solvers.AmericanTimeStepper import AmericanTimeStepper
from solvers.BermudamTimeStepper import BermudamTimeStepper

import unittest


class OptionsTestSuite(unittest.TestCase):
    """ Initialize all variables neede for the test. """

    def setUp(self):
        """ I don't know how to avoid this """

        self.m_r = float(0.05)
        self.m_q = 0.0
        self.m_sigma = float(0.4)
        self.m_exp_time = float(1.0)
        self.m_stock_price = 40
        self.m_strike = 50

        (self.m_ntimestep_p1, self.m_nsegs_p1) = (200, 200)
        (self.m_ntimestep_p2, self.m_nsegs_p2) = (200, 13)

        (self.m_smin, self.m_smax) = (10, 130)

        self.m_plot = True
        self.m_keep_hist = True
        self.m_error_tol = 0.12
        self.m_price_engine = None
        self.m_lagrange_p1 = LagrangeP1()
        self.m_lagrange_p2 = LagrangeP2()

    def _init_engine(self, option_style, shape_functions):
        """ Instantiate the right engine base on the parameters. """

        nsegs = {'P1': self.m_nsegs_p1,
                 'P2': self.m_nsegs_p2}[shape_functions]

        ntime_steps = {'P1': self.m_ntimestep_p1,
                       'P2': self.m_ntimestep_p2}[shape_functions]

        if 'E' == option_style:
            self.m_price_engine = nOp.NumericalOptionGBM(
                'E',
                self.m_exp_time,
                self.m_stock_price,
                self.m_strike,
                ntime_steps,
                nsegs,
                self.m_smin,
                self.m_smax,
                )
        elif 'A' == option_style:
            self.m_price_engine = nOp.NumericalOptionGBM(
                'A',
                self.m_exp_time,
                self.m_stock_price,
                self.m_strike,
                ntime_steps,
                nsegs,
                self.m_smin,
                self.m_smax,
                )
        else:
            assert False

        is_real = True
        self.m_price_engine.set_option_cos(
            self.m_sigma,
            self.m_q,
            self.m_r,
            is_real,
            is_real,
            is_real,
            )

        # self.m_price_engine.set_keep_hist(self.m_keep_hist)

        if 'P1' == shape_functions:
            self.m_price_engine.set_shape_functions(self.m_lagrange_p1)
        elif 'P2' == shape_functions:

            self.m_price_engine.set_shape_functions(self.m_lagrange_p2)

    @classmethod
    def __init_plotter(cls, style, call_type, va_di):
        """ Let plotter know what kind of data he'll be dealing with. """


        title = {'E': 'European', 'A': 'American',
                 'B': 'Bermudam'}[style]
        title += ' ' + {'P': 'Put', 'C': 'Call'}[call_type] + ' ' \
            + va_di

        styles = {'E': 'Analytical', 'A': 'Euro-Analytical',
                  'B': 'Euro-Analytical'}[style]

        plt_engine = DerivativePlots()
        plt_engine.set_title(title)
        plt_engine.set_curve_leyends(['Numeric', styles, 'Payoff'])

        plt_engine.set_axis_leyend({'x': 'Stoke Price', 'y': 'Option Price'
                               })
        return plt_engine

    def plot_greeks(self, stock_range, greeks, style, call_type):
        """ Method for greeks. """

        plt_engine = self.__init_plotter(style, call_type, 'vanilla')

        if 'P' == call_type:
            init_condition = clip(self.m_strike - stock_range, 0, 1e600)
            call_vanilla_exact = vanillaPutPrice(stock_range,
                    self.m_strike, self.m_r, self.m_sigma,
                    self.m_exp_time)
            plt_engine.set_leyend_posi(1)
        else:
            init_condition = clip(stock_range - self.m_strike, 0, 1e600)
            call_vanilla_exact = vanillaCallPrice(stock_range,
                    self.m_strike, self.m_r, self.m_sigma,
                    self.m_exp_time)
            plt_engine.set_leyend_posi(2)
        plt_engine.set_curves(stock_range, [greeks[:, 0], greeks[:, 1],
                          greeks[:, 2]])
        plt_engine.set_exit_time(1)
        plt_engine.plot_greeks()


    def plot_prices(self, stock_range, instru_value, style, call_type,
                    va_di='vanilla'):
        """ Method for plotting price values. """

        plt_engine = self.__init_plotter(style, call_type, va_di)

        if 'P' == call_type:
            init_condition = clip(self.m_strike - stock_range, 0, 1e600)
            if 'digital' == call_type:
                init_condition[init_condition > 0] = 1.0
            call_analytical = vanillaPutPrice(stock_range,
                    self.m_strike, self.m_r, self.m_sigma,
                    self.m_exp_time)
            plt_engine.set_leyend_posi(1)
        else:
            init_condition = clip(stock_range - self.m_strike, 0, 1e600)
            if 'vanilla' == va_di:
                call_analytical = vanillaCallPrice(stock_range,
                        self.m_strike, self.m_r, self.m_sigma,
                        self.m_exp_time)
                plt_engine.set_leyend_posi(2)

            if 'digital' == va_di:
                init_condition[init_condition > 0] = 1.0
                call_analytical = cashOrNothingCall(stock_range,
                        self.m_strike, self.m_r, self.m_sigma,
                        self.m_exp_time)
                plt_engine.set_leyend_posi(1)

        plt_engine.set_curves(stock_range, [instru_value, call_analytical,
                          init_condition])

        if True == self.m_keep_hist:
            price_hist = self.m_price_engine.get_price_history()
            plt_engine.set_history(stock_range, price_hist, self.m_exp_time)
            plt_engine.plot_history()
        else:
            plt_engine.plot()


class PlainVanillaEuroCall(OptionsTestSuite):
    """ Plain vanilla. You can play with all the parameters :) """

    def runTest(self):
        """ run regression test. """

        shape_functions = 'P2'
        self._init_engine('E', shape_functions)

        self.m_keep_hist = False
        time_stepper = EuropeanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_price_mode('greeks')
        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, greeks) = self.m_price_engine.price_call()

        if self.m_plot:
            self.plot_greeks(stock_range, greeks, 'E', 'C')

        call_vanilla_exact = vanillaCallPrice(stock_range,
                self.m_strike, self.m_r, self.m_sigma, self.m_exp_time)
        call_vanilla_exact[call_vanilla_exact < 1.0e-6] = 0.0
        error_norm = norm(greeks[:, 0] - call_vanilla_exact)
        assert self.m_error_tol > error_norm


class PlainVanillaEuroPut(OptionsTestSuite):
    """ Plain Put Euro style. """

    def runTest(self):
        """ run test. """

        shape_functions = 'P2'
        self._init_engine('E', shape_functions)

        time_stepper = EuropeanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, put_val) = self.m_price_engine.price_put()

        call_vanilla_exact = vanillaPutPrice(stock_range,
                self.m_strike, self.m_r, self.m_sigma, self.m_exp_time)

        if self.m_plot:
            self.plot_prices(stock_range, put_val, 'E', 'P')

        call_vanilla_exact[call_vanilla_exact < 1.0e-6] = 0.0
        error_norm = norm(put_val - call_vanilla_exact)
        assert self.m_error_tol > error_norm


class PlainVanillaAmerPut(OptionsTestSuite):
    """ American Put plain vanilla. """

    def runTest(self):
        """ run test. """

        shape_functions = 'P2'
        self._init_engine('A', shape_functions)

        time_stepper = AmericanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, put_val) = self.m_price_engine.price_put()

        if self.m_plot:
            self.plot_prices(stock_range, put_val, 'A', 'P')

        # no analytincal formula to compare.

        assert True


class PlainVanillaAmerCall(OptionsTestSuite):
    """ American call option. """

    def runTest(self):
        """ run test. """

        shape_functions = 'P2'
        self._init_engine('A', shape_functions)

        time_stepper = AmericanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2 + 50)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, call_val) = self.m_price_engine.price_call()

        if self.m_plot:
            self.plot_prices(stock_range, call_val, 'A', 'C')

        call_vanilla_exact = vanillaCallPrice(stock_range,
                self.m_strike, self.m_r, self.m_sigma, self.m_exp_time)
        call_vanilla_exact[call_vanilla_exact < 1.0e-6] = 0.0
        error_norm = norm(call_val - call_vanilla_exact)
        assert self.m_error_tol > error_norm


class BermudamCall(OptionsTestSuite):
    """ Bermudan call vanilla. """

    def runTest(self):
        """ run test. """

        shape_functions = 'P2'
        self._init_engine('E', shape_functions)

        exercise_dates = linspace(0.2, self.m_exp_time, 3, False)

        self.m_keep_hist = False
        time_stepper = BermudamTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_exercise_dates(exercise_dates)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, call_val) = self.m_price_engine.price_call()

        if self.m_plot:
            self.plot_prices(stock_range, call_val, 'B', 'C')

        call_value = vanillaCallPrice(stock_range, self.m_strike,
                self.m_r, self.m_sigma, self.m_exp_time)
        call_value[call_value < 1.0e-6] = 0.0
        error_norm = norm(call_val - call_value)
        assert self.m_error_tol > error_norm


class BermudamPut(OptionsTestSuite):
    """ Bermudam put vanilla. """

    def runTest(self):
        """ run test. """

        shape_functions = 'P2'
        self._init_engine('E', shape_functions)

        exercise_dates = linspace(0.2, self.m_exp_time, 3, False)

        self.m_keep_hist = False
        time_stepper = BermudamTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_exercise_dates(exercise_dates)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, call_val) = self.m_price_engine.price_put()

        if self.m_plot:
            self.plot_prices(stock_range, call_val, 'B', 'P')

        put_vanilla_value = vanillaPutPrice(stock_range, self.m_strike,
                self.m_r, self.m_sigma, self.m_exp_time)
        put_vanilla_value[put_vanilla_value < 1.0e-6] = 0.0

        # no analytincal formulat to compare.

        assert True


class DigitalEuroCall(OptionsTestSuite):

    """ Numerical digital functions. """

    def runTest(self):
        """ follow this steps to call the engine. """

        shape_functions = 'P2'
        self._init_engine('E', shape_functions)

        self.m_keep_hist = False
        time_stepper = EuropeanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p2)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_contract_type('digital')
        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, call_val) = self.m_price_engine.price_call()

        if self.m_plot:
            self.plot_prices(stock_range, call_val, 'E', 'C', 'digital')

        call_digital_value = cashOrNothingCall(stock_range,
                self.m_strike, self.m_r, self.m_sigma, self.m_exp_time)
        call_digital_value[call_digital_value < 1.0e-6] = 0.0
        error_norm = norm(call_val - call_digital_value)
        assert self.m_error_tol > error_norm


class PutCallParity(OptionsTestSuite):
    """ Consitency test. """

    def runTest(self):
        """ test put call parity. """

        shape_functions = 'P1'
        self._init_engine('E', shape_functions)

        time_stepper = EuropeanTimeStepper()
        time_stepper.set_number_steps(self.m_ntimestep_p1)
        time_stepper.set_time2simulate(self.m_exp_time)
        time_stepper.set_dolog(True)
        time_stepper.set_keep_hist(self.m_keep_hist)

        self.m_price_engine.set_time_stepper(time_stepper)

        (stock_range, put_val) = self.m_price_engine.price_put()
        (stock_range, call_val) = self.m_price_engine.price_call()

        error = call_val - put_val - stock_range + exp(-self.m_r
                                * self.m_exp_time) * self.m_strike
        assert self.m_error_tol > max(abs(error))


def my_suite():
    """ add new regression test here. """

    suite = unittest.TestSuite()

    suite.addTest(BermudamPut())
    suite.addTest(BermudamCall())
    suite.addTest(PutCallParity())
    suite.addTest(PlainVanillaEuroCall())
    suite.addTest(DigitalEuroCall())
    suite.addTest(PlainVanillaEuroPut())
    suite.addTest(PlainVanillaAmerCall())
    suite.addTest(PlainVanillaAmerPut())

    return suite


if __name__ == '__main__':

    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(my_suite())

